from walrus import *

# 连接到本地Redis服务器
db = Database(host='10.8.0.17', port=6379, db=4)

# 定义公司模型
class Company(Model):
    __database__ = db
    __namespace__ = 'company'
    corp_code = TextField(primary_key=True)
    users = HashField()          # 存储用户的详细信息，如DingTalk ID
    user_names = ZSetField()     # 用于分页查询的有序集合

    @classmethod
    # 添加或更新用户
    def add_or_update_user(self, name, dingtalk_id):
        self.users[name] = dingtalk_id
        self.user_names.add(0, name)  # 分数设置为0，仅排序用户名称

    # 分页查询公司的用户
    def fetch_users_page(self, page_number, page_size):
        start_index = (page_number - 1) * page_size
        end_index = start_index + page_size - 1
        user_names = self.user_names.range(start_index, end_index)
        users = [{'name': name, 'dingtalk_id': self.users[name]} for name in user_names]
        return users

    # 根据用户名查询用户
    def find_user(self, name):
        if name in self.users:
            return {'name': name, 'dingtalk_id': self.users[name]}
        return None
